import torch_npu
import torch
import subprocess
import shutil
from backends.backend_registry import register_backend, Backend
from utils.ascend_compile_pipeline import ascend_compile
from utils.correctness import execute_template
from utils.performance import time_execution_event_template
from config import project_root_path,ascendc_device
import os
import copy

@register_backend('ascendc')
class AscendBackend(Backend):
    def __init__(self):
        self.context = {}
        self.device = self.get_device()
    def get_device(self):
        return torch.device('npu:0')

    def get_hardware_name(self):
        return ascendc_device  # torch_npu.npu.get_device_name(device) causes crash

    def compile(self, generated_code, op, ref_src=""):
        try:
            info = ascend_compile(generated_code, op, self.context)
            return True, None, info
        except Exception as e:
            print(f"[DEBUG] Compile failed: {str(e)}")
            os.chdir(project_root_path)
            return False, str(e), None

    def correctness_execution(self, ref_src, context=None):
        synchronize = torch_npu.npu.synchronize
        try:
            exec(ref_src, context if context is not None else self.context)
        except Exception as e:
            raise RuntimeError(f"Failed to compile reference model: {str(e)}")
        return execute_template(synchronize, self.device, context if context is not None else self.context)

    def time_execution(self, eval_target='ModelNew', context=None):
        event_class = torch_npu.npu.Event
        synchronize = torch_npu.npu.synchronize
        return time_execution_event_template(context if context is not None else self.context, self.device, synchronize, event_class, eval_target)

    def cleanup(self, context=None, info=None):
        # Get operator identifier from context
        context = context if context is not None else self.context
        op_identifier = info.get('op_identifier') if info else None
        if op_identifier:
            # Uninstall the Python package
            try:
                subprocess.run(['pip', 'uninstall', op_identifier, '-y'], 
                             capture_output=True, text=True)
                print(f"[INFO] Uninstalled package: {op_identifier}")
            except Exception as e:
                print(f"[WARNING] Failed to uninstall package {op_identifier}: {e}")
            
            # Remove the temporary directory
            tmp_dir = os.path.join(project_root_path, 'tmp')
            op_dir = os.path.join(tmp_dir, op_identifier)
            if os.path.exists(op_dir):
                try:
                    shutil.rmtree(op_dir)
                    print(f"[INFO] Removed directory: {op_dir}")
                except Exception as e:
                    print(f"[WARNING] Failed to remove directory {op_dir}: {e}")
        
        # 强制清理NPU缓存和同步，确保资源完全释放
        torch_npu.npu.empty_cache()
        torch_npu.npu.synchronize(device=self.device)
        
        # 清理环境变量，避免影响后续kernel
        os.environ.pop("ASCEND_CUSTOM_OPP_PATH", None)
        
        # 安全地清理context
        if hasattr(context, 'context'):
            del context.context
